package com.gitee.wsl.math.random

import kotlin.math.pow
import kotlin.random.Random

fun <T> List<T>.randomWithWeights(weights: List<Double>, random: Random = Random): T = random.weighted(this.zip(weights).toMap())

fun <T> Random.weighted(weights: Map<T, Double>): T = shuffledWeighted(weights).first()
fun <T> Random.weighted(weights: RandomWeights<T>): T = shuffledWeighted(weights).first()

fun <T> Random.weighted(values: List<T>, weights: List<Double>): T = shuffledWeighted(values,weights).first()

fun <T> Random.shuffledWeighted(weights: Map<T, Double>): List<T> = shuffledWeighted(RandomWeights(weights))
fun <T> Random.shuffledWeighted(values: List<T>, weights: List<Double>): List<T> = shuffledWeighted(RandomWeights(values, weights))
fun <T> Random.shuffledWeighted(weights: RandomWeights<T>): List<T> {
    val randoms = (0 until weights.items.size).map { -(nextDouble().pow(1.0 / weights.normalizedWeights[it])) }
    val sortedIndices = (0 until weights.items.size).sortedWith { a, b -> randoms[a].compareTo(randoms[b]) }
    return sortedIndices.map { weights.items[it] }
}

data class RandomWeights<T>(val items: List<T>, val weights: List<Double>) {
    constructor(vararg pairs: Pair<T, Double>) : this(mapOf(*pairs))
    constructor(weightsMap: Map<T, Double>) : this(weightsMap.keys.toList(),weightsMap.values.toList())


    //val items = weightsMap.keys.toList()
    //val weights = weightsMap.values.toList()
    val normalizedWeights = normalizeWeights(weights)

    companion object {
        private fun normalizeWeights(weights: List<Double>): List<Double> {
            val min = weights.minOrNull() ?: 0.0
            return weights.map { (it + min) + 1 }
        }
    }
}